﻿#include "net.hpp"
#include "paramsparser.hpp"

#include <iostream>

using namespace leak;

#define MB / (1024 * 1024)

std::atomic_uint64_t g_connected{ 0 };
std::atomic_uint64_t g_qps{ 0 };
std::atomic_uint64_t g_recv{ 0 };

int         g_package = 16;
char*       g_data    = new char[g_package];
std::string g_host;
uint16_t    g_port = 0;

void EchoServer(utils::ContextPool pool) {
    net::tcp::Server::New(pool->Next())
        ->set_context_pool(pool)
        ->set_onaccept_error([](std::string_view error) {
            std::cerr << error << std::endl;
        })
        ->Listen("0.0.0.0", g_port)
        ->Accept([](std::shared_ptr<net::tcp::Socket> socket) {
            socket
                ->set_onerror([](std::string_view error) {
                    std::cout << "Server socket err: " << error << std::endl;
                })
                ->StartRead([](std::shared_ptr<net::tcp::Socket> socket, std::string_view data) {
                    auto bytes = data.length();
                    if (bytes < g_package) return 0;
                    socket->Write(data.substr(0, g_package));

                    return g_package;
                });
        });
}

void EchoClient(utils::ContextPool pool, size_t count) {
    for (int i = 0; i < count; i++) {
        net::tcp::Connect(
            pool->Next(),
            g_host,
            g_port,
            [](std::shared_ptr<net::tcp::Socket> socket) {
                g_connected++;
                socket
                    ->set_onerror([socket](std::string_view error) {
                        std::cout << "Client socket err: " << error << std::endl;
                    })
                    ->PrepareWrite(std::string_view(g_data, g_package))
                    ->Flush()
                    ->StartRead([](std::shared_ptr<net::tcp::Socket> socket, std::string_view data) {
                        auto bytes = data.length();
                        if (bytes < g_package) return 0;
                        g_qps++;
                        g_recv += g_package;
                        socket->Write(data.substr(0, g_package));
                        return g_package;
                    });
            },
            [](std::string_view err) {
                std::cout << "Client socket connect err: " << err << std::endl;
            });
    }
}

int main(int argc, char* argv[]) {
    ParamsParser parser(argc, argv);

    auto disable_server   = parser.GetParams<bool>("disable_server", false);
    auto disable_client   = parser.GetParams<bool>("disable_client", false);
    auto client_pool_size = parser.GetParams<int>("client_pool_size", 1);
    auto server_pool_size = parser.GetParams<int>("server_pool_size", 1);
    auto client_count     = parser.GetParams<int>("client_count", 1);
    auto url              = parser.GetParams<std::string_view>("url", "127.0.0.1:3050");

    g_package = parser.GetParams<int>("package", g_package);

    auto urlpos = url.find(':');
    g_host.assign(url.substr(0, urlpos));
    g_port = std::atoi(url.substr(urlpos + 1).data());

    std::cout << client_pool_size << " | " << server_pool_size << " | " << client_count << " | " << g_package << std::endl;

    if (!disable_server)
        EchoServer(std::make_shared<utils::SequenceContextPool>(server_pool_size));

    if (!disable_client) {
        auto pool = std::make_shared<utils::SequenceContextPool>(client_pool_size);
        EchoClient(pool, client_count);
        uint64_t last_recv = 0;
        uint64_t last_qps  = 0;
        while (true) {
            uint64_t recv = g_recv - last_recv;
            uint64_t qps  = g_qps - last_qps;

            std::cout << "connected: " << g_connected
                      << " | qps(W): " << (float)qps / 10000
                      << " | recv(MB): " << (float)recv MB
                      << std::endl;

            last_recv = g_recv;
            last_qps  = g_qps;
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }
    } else {
        getchar();
    }

    return 0;
}